import requests
from fastapi import APIRouter
from pydantic import BaseModel

from core.config import settings

class QueryRequest(BaseModel):
    text: str
    session_id: str

router = APIRouter(prefix="/rag", tags=["rag"])

@router.post("/query")
async def query_rag(query: QueryRequest):
    """
    调用rag接口
    """
    data = {
        "text": query.text,
        "session_id": query.session_id
    }
    res = requests.post(settings.RAG_HOST+'/api/v1/query', json=data)
    # log.info(res)
    if res.status_code == 200:
        return res.json()["answer"]
    else:
        return  None